Base

app/app.py

@@ -0,0 +1,28 @@
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+from .utils import recover_light_sources
+
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model=torch.load('model/model_epoch_49.pth',map_location=device)
+
+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(mask,input_image)
+    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()

app/utils.py

@@ -0,0 +1,61 @@
+import torch
+from PIL import Image
+import numpy as np
+
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model=torch.load('model/Unet_segformer.pth',map_location=device)
+
+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()
+    return mask
+
+def calculate_input_illuminance(image):
+    """
+    Calculate illuminance: I_input = C_r + C_g + C_b
+    """
+    return np.sum(image, axis=2)
+
+def generate_recovery_weight_matrix(illuminance_matrix, alpha=15):
+    """
+    Generate recovery weights using power function
+    Formula: W_r = ((I_input - min) / (max - min))^α
+    """
+    I_min = np.min(illuminance_matrix)
+    I_max = np.max(illuminance_matrix)
+    
+    if I_max == I_min:
+        normalized = np.zeros_like(illuminance_matrix)
+    else:
+        normalized = (illuminance_matrix - I_min) / (I_max - I_min)
+    
+    # Apply power function with α = 15
+    W_r = np.power(normalized, alpha)
+    return W_r
+
+def recover_light_sources(original_image, network_output, alpha=15):
+    """
+    Final recovery: I_final = (1 - W_r) ⊙ N(C) + W_r ⊙ C
+    """
+    # Calculate illuminance and recovery weights
+    I_input = calculate_input_illuminance(original_image)
+    W_r = generate_recovery_weight_matrix(I_input, alpha)
+    
+    # Expand to match image dimensions
+    W_r_expanded = np.expand_dims(W_r, axis=2)
+    W_r_expanded = np.repeat(W_r_expanded, 3, axis=2)
+    
+    # Convex combination for light source recovery
+    I_final = (1 - W_r_expanded) * network_output + W_r_expanded * original_image
+    return np.clip(I_final, 0, 255).astype(np.uint8)

main.py

@@ -0,0 +1,30 @@
+from fastapi import FastAPI, File, UploadFile
+from pydantic import BaseModel
+from utils import predict
+from starlette.responses import StreamingResponse, Response
+from io import BytesIO
+from PIL import Image
+app=FastAPI()
+
+
+@app.get("/")
+def home():
+	return {"health_check":"OK"} 
+	
+#@app.post("/remove")
+#async def remove(file: UploadFile = File(...)):
+#	contents=await file.read()
+#	image = Image.open(BytesIO(contents)).convert("RGB")
+#	mask=predict(image)
+#	return {"mask": mask.tolist()}
+
+@app.post("/remove")
+async def remove(file: UploadFile = File(...)):
+	contents=await file.read()
+	image = Image.open(BytesIO(contents)).convert("RGB")
+	mask=predict(image)
+	img=Image.fromarray((mask*255).astype('uint8'))
+	buffered=BytesIO()
+	img.save(buffered,format="JPEG")
+	img_str=buffered.getvalue()
+	return Response(content=img_str, media_type="image/jpeg")

model/model_epoch_49.pth

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