Update app.py

app.py

@@ -11,26 +11,38 @@ import numpy as np
 class ZeroDCE(nn.Module):
     def __init__(self):
         super(ZeroDCE, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
-        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
-        self.conv3 = nn.Conv2d(32, 32, 3, padding=1)
-        self.conv4 = nn.Conv2d(32, 24, 3, padding=1)
         self.relu = nn.ReLU(inplace=True)
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1, bias=True)
+        self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
+        self.conv3 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
+        self.conv4 = nn.Conv2d(32, 24, 3, padding=1, bias=True)
+        self.conv5 = nn.Conv2d(24, 8, 3, padding=1, bias=True)
 
     def forward(self, x):
         x1 = self.relu(self.conv1(x))
         x2 = self.relu(self.conv2(x1))
         x3 = self.relu(self.conv3(x2))
-        x_r = torch.tanh(self.conv4(x3))
-        return x_r
+        x4 = self.relu(self.conv4(x3))
+        out = torch.tanh(self.conv5(x4))
+        return out
 
 def enhance_image(img, model):
-    img_tensor = torch.from_numpy(np.array(img)).float() / 255.0
-    img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
+    # Convert PIL -> Tensor
+    img_np = np.array(img).astype(np.float32) / 255.0
+    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
+
     with torch.no_grad():
-        enhanced = model(img_tensor)
+        enhancement_map = model(img_tensor)
+
+    # Apply enhancement: simple iterative curve
+    enhanced = img_tensor
+    for i in range(8):
+        enhanced = enhanced + enhancement_map[:, i*3:(i+1)*3, :, :] * (enhanced**2 - enhanced)
+
+    enhanced = torch.clamp(enhanced, 0, 1)
     enhanced = enhanced.squeeze(0).permute(1, 2, 0).cpu().numpy()
-    enhanced = np.clip(enhanced * 255, 0, 255).astype(np.uint8)
+    enhanced = (enhanced * 255).astype(np.uint8)
+
     return Image.fromarray(enhanced)
 
 # -----------------------------