uploader/nice.py

#!/usr/bin/env python3
from collections import defaultdict
from pathlib import Path

from bs4 import BeautifulSoup
from local_file_picker import local_file_picker

from nicegui import ui, app, html, run

from datetime import datetime

import zmq
import zmq.asyncio

from loguru import logger
import sys

from anonymiser import Anonymizer
import pydicom

import requests
import time

from multiprocessing import Manager, Queue, freeze_support

import watchfiles


from typing import Optional

from fastapi import Request
from fastapi.responses import RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware
#import typer

#typer_app = typer.Typer()

#logger.add(sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO")

#async def pick_file() -> None:
#    result = await local_file_picker('~', multiple=True)
#    ui.notify(f'You chose {result}')
#    ui.html("<h2>Uploader tool</h2>")
#
#    ui.chat_message(f'{result}!',
#                name='Robot',
#                stamp='now',
#                avatar='https://robohash.org/ui')
#
#
#@ui.page('/')
#def index():
#    ui.html("<h2>Uploader tool</h2>")
#    ui.button('Choose file', on_click=pick_file, icon='folder')

EXPORT_PATH = Path("./test/export")
PROCESS_PATH = Path("./test/to_process")
ANON_PATH = Path("./test/anon")

EXPORT_PATH.mkdir(exist_ok=True, parents=True)
PROCESS_PATH.mkdir(exist_ok=True, parents=True)
EXPORT_PATH.mkdir(exist_ok=True, parents=True)

BASE_SITE_URL = "http://localhost:8000"

LOGIN_URL = f"{BASE_SITE_URL}/accounts/login/"

LOGIN_SUCCESS = False

rqst = None


use_zmq = False
loaded_files: dict = {}
uploaded_files: dict = {}
duplicate_series = set()

if use_zmq:
    context = zmq.asyncio.Context()
    socket = context.socket(zmq.PULL)
    #socket.bind('tcp://localhost:5575')
    try:
        socket.bind('tcp://localhost:9975')
    except zmq.error.ZMQError as e:
        logger.debug("send reload")
        socket = context.socket(zmq.PUSH)
        socket.connect('tcp://localhost:5575')
        socket.send_string("loaded")
        sys.exit(0)

    poller = zmq.asyncio.Poller()
    poller.register(socket, zmq.POLLIN)


anonymizer = Anonymizer(ANON_PATH)

def load_files(q: Queue):
    loaded_files: dict = {}
    loaded_series = defaultdict(list)
    loaded_series_data = {}
    # Move files
    src_path = EXPORT_PATH

    to_process = []
    for file in src_path.iterdir():
        logger.debug(f"Moving {file} to {PROCESS_PATH}")

        file_to_process = PROCESS_PATH / file.name
        file.rename(file_to_process)

        to_process.append(file_to_process)

    to_process_len = len(to_process)
    for n, file in enumerate(to_process):
        dataset, output_file = anonymizer.anonymize_file(file, remove_original=True)
        #with processed_files:
        #    ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {file.name} -> {output_file.name}")

        loaded_files[output_file] = (dataset.StudyDescription, dataset.SeriesDescription, str(dataset.SeriesInstanceUID))

        loaded_series[dataset.SeriesInstanceUID].append(output_file)
        loaded_series_data[dataset.SeriesInstanceUID] = (dataset.StudyDescription, dataset.SeriesDescription)

        q.put_nowait((n + 1) / to_process_len)

        #time.sleep(0.01)

    return loaded_files, loaded_series, loaded_series_data 
    #load_series_view()



def reload_anonymized(q: Queue):
    loaded_files = {}
    loaded_series = defaultdict(list)
    loaded_series_data = {}

    #processed_files.clear()


    to_process = list(ANON_PATH.glob("*.dcm"))
    to_process_len = len(to_process)
    for n, file in enumerate(to_process):
        with pydicom.dcmread(file) as dataset:
            loaded_files[file] = (dataset.StudyDescription, dataset.SeriesDescription, str(dataset.SeriesInstanceUID))

            loaded_series[dataset.SeriesInstanceUID].append(file)
            loaded_series_data[dataset.SeriesInstanceUID] = (dataset.StudyDescription, dataset.SeriesDescription)

        #with processed_files:
        #    ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {file.name} -> Reloaded")


        q.put_nowait(f"{(n + 1) / to_process_len * 100:.0f}%")
        logger.debug(file)

        #time.sleep(0.01)
    
    logger.debug("end reload")
    logger.debug(loaded_files)
    logger.debug(loaded_series)
    logger.debug("-----")

    return loaded_files, loaded_series, loaded_series_data 

    #load_series_view()

def upload_files(q, rqst):
    #global rqst#, uploaded_files

    uploaded_files = {}

    files_to_upload = []
    for file in ANON_PATH.iterdir():
        files_to_upload.append(("files", open(str(file), "rb")))

    if not files_to_upload:
        ui.notify('No files to upload', color='negative')
        return

    # chunck files
    n = 10
    chunked_files = [
        files_to_upload[i : i + n] for i in range(0, len(files_to_upload), n)
    ]


    upload_file_list = []
    duplicate_file_list = []
    failed = []
    duplicate_series = set()

    logger.debug("START UPLOAD")
    logger.debug(files_to_upload)

    to_process_len = len(chunked_files)

    for n, files in enumerate(chunked_files):

        def upload_files_(files):
            rqst.headers["X-CSRFToken"] = rqst.cookies["csrftoken"]
            # print(self.rqst.headers)
            # print(self.rqst.cookies)
            resp = rqst.post(
                f"{BASE_SITE_URL}/api/atlas/upload_dicom",
                # data=data,
                files=files,
            )

            logger.debug(f"Resp: {resp}")
            logger.debug(f"{resp.content}")
            return resp

        # progress_dialog.Update(n, f"Uploading batch {n}/{len(chunked_files)}")

        logger.debug(f"n: {n}")
        # try to upload the files

        for i in range(3):
            logger.error(f"i: {i}")
            resp = upload_files_(files)
            if resp.status_code == 200:
                upload_file_list.extend(resp.json()["uploaded"])
                duplicate_file_list.extend(resp.json()["duplicates"])
                failed.extend(resp.json()["failed"])
                duplicate_series.update(resp.json()["duplicate_series"])

                break

            logger.debug(f"n: {n} fail (attempt {i})")

        q.put_nowait((n + 1) / to_process_len)
        # progress_dialog.Destroy()
    print(upload_file_list)
    print("dup", duplicate_file_list)
    print("failed", failed)

    for f, hash in upload_file_list:
        uploaded_files[f] = "success"

    for f, hash in duplicate_file_list:
        uploaded_files[f] = "duplicate"

    for f in failed:
        uploaded_files[f] = "failed"


    return uploaded_files, upload_file_list, duplicate_file_list, failed, duplicate_series







@ui.page('/login')
def login() -> Optional[RedirectResponse]:
    dark = ui.dark_mode()
    dark.enable()
    def try_login() -> None:  # local function to avoid passing username and password as arguments
        logger.debug(username.value)
        #if passwords.get(username.value) == password.value:
        #    app.storage.user.update({'username': username.value, 'authenticated': True})
        #    ui.navigate.to(app.storage.user.get('referrer_path', '/'))  # go back to where the user wanted to go
        #else:
        #    ui.notify('Wrong username or password', color='negative')
        user = username.value
        pw = password.value

        global rqst, LOGIN_SUCCESS
        rqst = requests.session()
        try: 
            rsp = rqst.get(LOGIN_URL)
        except requests.exceptions.ConnectionError:
            ui.notify('Connection error!', color='negative')
            return

        token = (
            BeautifulSoup(rsp.content, "html.parser").find("input").attrs["value"]
        )  # , attr={"name": "csrfmiddlewaretoken"}).attrs("value")

        # token = rsp.cookies["csrftoken"]
        # header = {"X-CSRFToken": token}
        # cookies = {"csrftoken": token}

        data = {
            "username": user,
            "password": pw,
            "csrfmiddlewaretoken": token,
            "next": "/",
        }

        rqst.headers["Referer"] = LOGIN_URL

        rsp = rqst.post(
            LOGIN_URL,
            data=data,
            # headers=header,
            # cookies=cookies
        )

        soup = BeautifulSoup(rsp.content, "html.parser")

        if soup.find("button") and soup.find("button").find(string="Login"):
            print("login fail")
            pass
        else:
            print("login success")
            LOGIN_SUCCESS = True
            user_label.clear()
            user_label.text = f"User: {user}"
        # print(rqst.headers)

        rqst.headers["X-CSRFToken"] = token

        print(rqst.headers)
        print(rqst.cookies)
        new = rqst.get(f"{BASE_SITE_URL}/accounts/api/atlas/get_cases_user")


        if LOGIN_SUCCESS:
            ui.notify('Logged in!', color='positive')
            ui.navigate.to("/")
            #login_button.visible = False
            pass
            #self.parent.upload_button.Enable()
            #self.parent.login_button.Hide()
            #self.parent.panel.Layout()
            #self.Destroy()

            # self.parent.login_button.SetBackgroundColour(wx.BLUE)
            # self.parent.login_button.SetForegroundColour(wx.WHITE)
        else:
            ui.notify('Wrong username or password', color='negative')
            pass
            #self.login_fail()

    #if app.storage.user.get('authenticated', False):
    #    return RedirectResponse('/')
    with ui.card().classes('absolute-center'):
        username = ui.input('Username').on('keydown.enter', try_login)
        password = ui.input('Password', password=True, password_toggle_button=True).on('keydown.enter', try_login)
        ui.button('Log in', on_click=try_login)
    return None

@ui.page('/')
def main_page():
    def load_series_view(loaded_series, loaded_series_data):
        logger.debug("load series view")
        logger.debug(loaded_series)
        for key in loaded_series:
            logger.debug(f"load series: {key}")
            with series:
                with ui.card():
                    ui.label(loaded_series_data[key][0])
                    ui.label(loaded_series_data[key][1])
                    ui.label(f"Images: {len(loaded_series[key])}")

        series_title.visible = True

    if use_zmq:
        async def read_loop() -> None:
            while not app.is_stopped:
                events = await poller.poll()
                if socket in dict(events):
                    data = await socket.recv()
                    msg = data.decode()

                    logger.debug(f'Received msg {msg}')

                    match msg:
                        case "loaded":
                            with updates:
                                ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {msg}")
                            logger.debug("loaded")
                            load_files_start()

                            
                        case _:
                            logger.error("invalid message")
        app.on_startup(read_loop)

    dark = ui.dark_mode()
    dark.enable()
    ui.page_title('Uploader tool')
    ui.label('Uploader tool').classes("text-h2")
    user_label = ui.label("User: None")

    login_button = ui.button("LOGIN", on_click=lambda: ui.navigate.to("/login"))
    login_button.bind_visibility_from(globals(), "LOGIN_SUCCESS", backward=lambda v: not v)

    updates = ui.list()

    series_title = ui.label('Series to upload').classes("text-h3")
    series_title.visible = False

    series = ui.row()



    async def load_files_start():
        global loaded_files
        anon_progress.visible = True
        loaded_files, loaded_series, loaded_series_data = await run.cpu_bound(load_files, queue)
        anon_progress.visible = False
        load_series_view(loaded_series, loaded_series_data)
        loaded_files_ui.refresh()

    async def reload_anonymized_start():
        global loaded_files
        anon_progress.visible = True
        loaded_files, loaded_series, loaded_series_data = await run.cpu_bound(reload_anonymized, queue)
        anon_progress.visible = False
        load_series_view(loaded_series, loaded_series_data)
        loaded_files_ui.refresh()

    async def upload_files_start():
        global uploaded_files, duplicate_series
        upload_progress.visible = True
        uploaded_files, upload_file_list, duplicate_file_list, failed, duplicate_series = await run.cpu_bound(upload_files, upload_queue, rqst)
        upload_progress.visible = False

        ui.notify(f"Uploaded {len(upload_file_list)} files", color='positive')
        ui.notify(f"Duplicate {len(duplicate_file_list)} files", color='warning')
        ui.notify(f"Failed {len(failed)} files", color='negative')

        uploaded_files_ui.refresh()

    upload_button = ui.button("Upload", on_click=upload_files_start)
    upload_button.bind_visibility_from(globals(), "LOGIN_SUCCESS")
#    typer_app()
    # Create a queue to communicate with the heavy computation process
    queue = Manager().Queue()
    # Update the progress bar on the main process
    ui.timer(0.1, callback=lambda: progressbar.set_value(queue.get() if not queue.empty() else progressbar.value))

    # Create a queue to communicate with the heavy computation process
    upload_queue = Manager().Queue()
    # Update the progress bar on the main process
    ui.timer(0.1, callback=lambda: upload_progressbar.set_value(upload_queue.get() if not upload_queue.empty() else upload_progressbar.value))

    with ui.expansion('Extra', icon='build').classes('w-full'):
        ui.button("Load files", on_click=load_files_start)
        ui.button("Reload anon", on_click=reload_anonymized_start)



    upload_progress = ui.row().classes("w-full place-content-center bg-blue-900")

    with upload_progress:
        ui.label("Uploading files")
        upload_progressbar = ui.linear_progress(value=0).props('instant-feedback')
    upload_progress.visible = False

    anon_progress = ui.row().classes("w-full place-content-center bg-blue-900")

    with anon_progress:
        ui.label("Anonymizing files")
        progressbar = ui.linear_progress(value=0).props('instant-feedback')
    anon_progress.visible = False

    @ui.refreshable
    def loaded_files_ui() -> None:
        with ui.expansion('Loaded files', icon='drive_folder_upload').classes('w-full'):
            ui.label(f"Items: {len(loaded_files)}")
            with ui.list():
                for file, data in loaded_files.items():
                    ui.item(f"{file} - {data[0]} - {data[1]}")

    @ui.refreshable
    def uploaded_files_ui() -> None:
        with ui.expansion('Uploaded files', icon='file_upload').classes('w-full'):
            ui.label(f"Items: {len(uploaded_files)}")
            with ui.list():
                for series in duplicate_series:
                    ui.item(f"Duplicate series: {series}", on_click=lambda: ui.navigate.to(f'{BASE_SITE_URL}{series}', new_tab=True)).classes("text-red-500")
            with ui.list():
                for file, data in uploaded_files.items():
                    ui.item(f"{file} - {data}")

    with ui.expansion('File status', icon='view_list').classes('w-full'):
        loaded_files_ui()
        uploaded_files_ui()

        #with ui.expansion('Processed files!', icon='work').classes('w-full'):
        #    processed_files = ui.list()

    with ui.expansion("Settings", icon="settings").classes("w-full"):
        ui.label("Settings").classes("text-h4")
        ui.label(f"Export path: {EXPORT_PATH}")
        ui.label(f"Process path: {PROCESS_PATH}")
        ui.label(f"Anon path: {ANON_PATH}")




if __name__ in {"__main__", "__mp_main__"}:
    freeze_support()

    ui.run(storage_secret="123456", reload=False)