BlurMe/blurme/state.py
2024-01-16 19:21:51 +01:00

95 lines
3.2 KiB
Python

"""Base state for the app."""
import reflex as rx
import os
import stat
#from blurme import styles
import asyncio
from graphics.image_modification import blur_boxes_on_image
from graphics.image_modification import show_image_with_boxes
from ml.element_detection import BoundBox, detect
class State(rx.State):
"""The app state."""
# The images to show.
img: list[str] = []
show: bool = False
def toggle_show(self):
self.show = not self.show
show_extension_alert: bool = False
def toggle_extension_alert(self):
self.show_extension_alert = not self.show_extension_alert
def change(self):
self.show = not (self.show)
async def handle_upload(self, files: list[rx.UploadFile]):
"""Handle the upload of file(s).
Args:
files: The uploaded files.
"""
max_files_limit = 5
acceptable_extensions = {"png", "jpg", "jpeg", "webp", "tif", "tiff"}
if len(self.img) + len(files) > max_files_limit:
self.toggle_show()
return
for file in files:
upload_data = await file.read()
outfile = rx.get_asset_path(file.filename)
extension = file.filename.split(".")[-1]
if extension.lower() not in acceptable_extensions:
self.toggle_extension_alert()
return
# Save the file.
with open(outfile, "wb") as file_object:
file_object.write(upload_data)
# Update the img var.
self.img.append(file.filename)
self.img = self.img
def delete_image(self, img_name: str):
if img_name in self.img:
self.img.remove(img_name)
self.img=self.img
def download_image(self, img_name):
print(self.img)
return rx.download(url=f'/{img_name}', filename=img_name)
async def image_analysis(self, img_name: str):
if img_name in self.img:
stripped_image_name = str(img_name).lstrip("analysis-").lstrip("anonim-")
image_path = rx.get_asset_path(stripped_image_name)
#timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
new_image_name = f"analysis-{stripped_image_name}"
new_image_path = rx.get_asset_path(new_image_name)
show_image_with_boxes(image_path, detect(image_path), new_image_path)
upload_file = rx.UploadFile(file=open(new_image_path, 'rb'), filename=new_image_name)
self.delete_image(img_name)
await self.handle_upload([upload_file])
self.img=self.img
async def image_anonymization(self, img_name: str):
if img_name in self.img:
stripped_image_name = str(img_name).lstrip("analysis-").lstrip("anonim-")
image_path = rx.get_asset_path(stripped_image_name)
#timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
new_image_name = f"anonim-{stripped_image_name}"
new_image_path = rx.get_asset_path(new_image_name)
blur_boxes_on_image(image_path, detect(image_path), new_image_path)
upload_file = rx.UploadFile(file=open(new_image_path, 'rb'), filename=new_image_name)
self.delete_image(img_name)
await self.handle_upload([upload_file])
self.img=self.img