Final fixes

.python-version

@@ -0,0 +1 @@
+3.13

app.py

@@ -4,55 +4,44 @@ from PIL import Image
 import numpy as np
 from app.utils import recover_light_sources
 from model.model import model
+from torchvision import transforms
+
 
 device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
 chk=torch.load('model/model_epoch_49.pth',map_location=device)
 model.load_state_dict(chk['model_state_dict'])
 
-#def evaluate(model,image):
-#    model.eval()
-#    with torch.no_grad():
-#        image = image.to(device)
-#        #outputs= model(image.unsqueeze(0))
-#        outputs= model(image)
-#        return outputs.squeeze(0).squeeze(0).cpu()
-    
-#def predict(input_image):
-    #input_image=Image.open(inp_img).convert('RGB')
-#    input_image=input_image.resize((512,512))
-#    input_image_torch=torch.tensor(np.array(input_image)).permute(2,0,1).unsqueeze(0).float()/255.0
-#    mask=evaluate(model,input_image_torch)
-#    mask=mask.permute(1,2,0).numpy()
-#    final_img=recover_light_sources(input_image,mask)
-#    return final_img
-
-def evaluate(model, image_tensor):
+transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor(),  # -> (C,H,W), dtype=float32, range [0,1]
+])
+
+def evaluate(model, image):
     """
     Run the model on the given image tensor and return output as numpy array (H,W,C) in [0,255].
     """
+    model.eval()
     with torch.no_grad():
-        image_tensor = image_tensor.to(device, dtype=torch.float32)
-        outputs = model(image_tensor)
+        #image = image.to(device, dtype=torch.float32)
+        outputs = model(image)
         outputs = torch.clamp(outputs, 0.0, 1.0)
         outputs_np = outputs.squeeze(0).permute(1, 2, 0).cpu().numpy()
-        return (outputs_np * 255).astype(np.uint8)  # Convert to uint8
+        return (outputs_np * 255).astype(np.uint8)  # Convert to uint8 for recovery step
 
 def predict(input_image):
     """
     Predict clean image from flare image, then recover light sources.
     """
-    # Ensure input image is in RGB
-    if input_image.mode != "RGB":
-        input_image = input_image.convert("RGB")
     # Resize and prepare input tensor
-    input_image_resized = input_image.resize((512, 512))
-    input_np = np.array(input_image_resized) / 255.0  # scale to [0, 1]
-    input_tensor = torch.from_numpy(input_np).permute(2, 0, 1).unsqueeze(0).float()
-    # Get predicted clean image from model
-    pred_clean_img = evaluate(model, input_tensor)
-    # Return as PIL Image
-    return Image.fromarray(pred_clean_img)
+    input_img = input_image.convert('RGB').resize((512, 512), Image.BILINEAR)
 
+    input_tensor = transform(input_img).unsqueeze(0).to(device, dtype=torch.float32)
+
+    # Get predicted clean image from model
+    pred_clean_img = evaluate(model, input_tensor)  # uint8 predicted clean
+    # Recover light sources
+    final_img = recover_light_sources(network_output=pred_clean_img,original_image=input_img)
+    return final_img
 
 demo = gr.Interface(fn=predict, inputs=gr.Image(type="pil"),outputs=gr.Image(), examples=["test_imgs/test1.png", "test_imgs/test2.png","test_imgs/test3.png"])
 demo.launch()

main.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from flare-removal-2-0!")
+
+
+if __name__ == "__main__":
+    main()

model/model_epoch_49.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d03f06fc9cf54ceaf33ad68481c41c3a281137f262d356898aadafcd6ec9ae85
-size 197469247
+oid sha256:6e1d2df83637ae40a4569e93f8135dcd248d49595d9cba685e16aa7f3aec798e
+size 197467007

pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "flare-removal-2-0"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "fastapi>=0.116.1",
+    "gradio>=5.44.1",
+    "numpy>=2.3.2",
+    "pillow>=11.3.0",
+    "pydantic>=2.11.7",
+    "segmentation-models-pytorch>=0.5.0",
+    "torch>=2.8.0",
+    "uvicorn>=0.35.0",
+]