Update app.py

app.py

@@ -1,162 +1,170 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.models as models
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-import numpy as np
-import cv2
-import base64
-from io import BytesIO
-from PIL import Image
-
-app = Flask(__name__)
-CORS(app)
-
-DEVICE = torch.device('cpu') # Force CPU for HF Free Tier
-MODEL_PATH = "best_model_fixed.pth" # Upload your trained .pth file
-
-class InpaintingGenerator(nn.Module):
-    def __init__(self, input_channels=4):
-        super().__init__()
-        resnet = models.resnet34(weights=None)
-        
-        self.enc1 = nn.Sequential(
-            nn.Conv2d(input_channels, 64, kernel_size=7, stride=2, padding=3, bias=False),
-            resnet.bn1, resnet.relu
-        )
-        self.enc2 = resnet.layer1
-        self.enc3 = resnet.layer2
-        self.enc4 = resnet.layer3
-        self.enc5 = resnet.layer4
-        
-        self.bottleneck = nn.Sequential(
-            nn.Conv2d(512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(True),
-            nn.Conv2d(512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(True)
-        )
-        
-        self.up1 = self._make_decoder_block(512, 256)
-        self.up2 = self._make_decoder_block(512, 128) # 256+256
-        self.up3 = self._make_decoder_block(256, 64)  # 128+128
-        self.up4 = self._make_decoder_block(128, 32)  # 64+64
-        
-        self.texture_refine = nn.Sequential(
-            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(True),
-            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(True)
-        )
-        
-        self.final = nn.Sequential(
-            nn.Conv2d(32, 16, 3, padding=1), nn.ReLU(True),
-            nn.Conv2d(16, 3, 3, padding=1), nn.Tanh()
-        )
-        
-    def _make_decoder_block(self, in_channels, out_channels):
-        return nn.Sequential(
-            nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1),
-            nn.BatchNorm2d(out_channels), nn.ReLU(True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1),
-            nn.BatchNorm2d(out_channels), nn.ReLU(True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1),
-            nn.BatchNorm2d(out_channels), nn.ReLU(True)
-        )
-
-    def forward(self, img, mask):
-        x = torch.cat([img, mask], dim=1)
-        x1 = self.enc1(x)
-        x2 = self.enc2(x1)
-        x3 = self.enc3(x2)
-        x4 = self.enc4(x3)
-        x5 = self.enc5(x4)
-        
-        x = self.bottleneck(x5)
-        
-        x = self.up1(x)
-        x = torch.cat([x, x4], dim=1)
-        x = self.up2(x)
-        x = torch.cat([x, x3], dim=1)
-        x = self.up3(x)
-        x = torch.cat([x, x2], dim=1)
-        x = self.up4(x)
-        
-        x = self.texture_refine(x)
-        return self.final(x)
-
-
-print("Loading Inpainting Model...")
-model = InpaintingGenerator().to(DEVICE)
-
-try:
-    # Set weights_only=False to avoid numpy errors
-    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
-    
-    # Handle DataParallel wrapping
-    if 'generator' in checkpoint:
-        state_dict = checkpoint['generator']
-    else:
-        state_dict = checkpoint
-        
-    new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
-    model.load_state_dict(new_state_dict, strict=False)
-    model.eval()
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-
-
-def to_base64(image_array):
-    img = Image.fromarray(image_array)
-    buffer = BytesIO()
-    img.save(buffer, format="PNG")
-    return base64.b64encode(buffer.getvalue()).decode('utf-8')
-
-@app.route('/')
-def home():
-    return "Inpainting API is Running!"
-
-@app.route('/inpaint', methods=['POST'])
-def inpaint():
-    if 'image' not in request.files or 'mask' not in request.files:
-        return jsonify({'error': 'Please upload both image and mask'}), 400
-    
-    try:
-        img_file = request.files['image']
-        img_arr = np.frombuffer(img_file.read(), np.uint8)
-        img_cv = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
-        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
-        
-        # 2. Read Mask
-        mask_file = request.files['mask']
-        mask_arr = np.frombuffer(mask_file.read(), np.uint8)
-        mask_cv = cv2.imdecode(mask_arr, cv2.IMREAD_GRAYSCALE)
-        
-        # 3. Preprocess
-        img_h, img_w = img_cv.shape[:2]
-        # Resize to 512x512 for model
-        img_resized = cv2.resize(img_cv, (512, 512))
-        mask_resized = cv2.resize(mask_cv, (512, 512))
-        
-        # Normalize
-        img_tensor = (torch.tensor(img_resized).float() / 127.5) - 1.0
-        img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).to(DEVICE)
-        
-        mask_tensor = (torch.tensor(mask_resized).float() > 127).float().unsqueeze(0).unsqueeze(0).to(DEVICE)
-        
-        # 4. Inference
-        with torch.no_grad():
-            output = model(img_tensor, mask_tensor)
-            
-        # 5. Post-process
-        output_np = output.squeeze().permute(1, 2, 0).cpu().numpy()
-        output_np = (output_np + 1.0) * 127.5
-        output_np = np.clip(output_np, 0, 255).astype(np.uint8)
-        
-        # Resize back to original dimensions
-        output_final = cv2.resize(output_np, (img_w, img_h))
-        
-        return jsonify({'result': f"data:image/png;base64,{to_base64(output_final)}"})
-        
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-if __name__ == '__main__':
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import numpy as np
+import cv2
+import base64
+from io import BytesIO
+from PIL import Image
+
+app = Flask(__name__)
+CORS(app)
+
+DEVICE = torch.device('cpu') # Force CPU for HF Free Tier
+MODEL_PATH = "best_model_fixed.pth" # Upload your trained .pth file
+
+class InpaintingGenerator(nn.Module):
+    def __init__(self, input_channels=4):
+        super().__init__()
+        resnet = models.resnet34(weights=None)
+        
+        self.enc1 = nn.Sequential(
+            nn.Conv2d(input_channels, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            resnet.bn1, resnet.relu
+        )
+        self.enc2 = resnet.layer1
+        self.enc3 = resnet.layer2
+        self.enc4 = resnet.layer3
+        self.enc5 = resnet.layer4
+        
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(True),
+            nn.Conv2d(512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(True)
+        )
+        
+        self.up1 = self._make_decoder_block(512, 256)
+        self.up2 = self._make_decoder_block(512, 128) # 256+256
+        self.up3 = self._make_decoder_block(256, 64)  # 128+128
+        self.up4 = self._make_decoder_block(128, 32)  # 64+64
+        
+        self.texture_refine = nn.Sequential(
+            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(True),
+            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(True)
+        )
+        
+        self.final = nn.Sequential(
+            nn.Conv2d(32, 16, 3, padding=1), nn.ReLU(True),
+            nn.Conv2d(16, 3, 3, padding=1), nn.Tanh()
+        )
+        
+    def _make_decoder_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels), nn.ReLU(True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            nn.BatchNorm2d(out_channels), nn.ReLU(True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            nn.BatchNorm2d(out_channels), nn.ReLU(True)
+        )
+
+    def forward(self, img, mask):
+        x = torch.cat([img, mask], dim=1)
+        x1 = self.enc1(x)
+        x2 = self.enc2(x1)
+        x3 = self.enc3(x2)
+        x4 = self.enc4(x3)
+        x5 = self.enc5(x4)
+        
+        x = self.bottleneck(x5)
+        
+        x = self.up1(x)
+        x = torch.cat([x, x4], dim=1)
+        x = self.up2(x)
+        x = torch.cat([x, x3], dim=1)
+        x = self.up3(x)
+        x = torch.cat([x, x2], dim=1)
+        x = self.up4(x)
+        
+        x = self.texture_refine(x)
+        return self.final(x)
+
+
+print("Loading Inpainting Model...")
+model = InpaintingGenerator().to(DEVICE)
+
+try:
+    # Set weights_only=False to avoid numpy errors
+    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
+    
+    # Handle DataParallel wrapping
+    if 'generator' in checkpoint:
+        state_dict = checkpoint['generator']
+    else:
+        state_dict = checkpoint
+        
+    new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+    model.load_state_dict(new_state_dict, strict=False)
+    model.eval()
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+
+
+def to_base64(image_array):
+    img = Image.fromarray(image_array)
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+@app.route('/')
+def home():
+    return "Inpainting API is Running!"
+
+@app.route('/inpaint', methods=['POST'])
+def inpaint():
+    if 'image' not in request.files or 'mask' not in request.files:
+        return jsonify({'error': 'Please upload both image and mask'}), 400
+    
+    try:
+        # 1. Read Image
+        img_file = request.files['image']
+        img_arr = np.frombuffer(img_file.read(), np.uint8)
+        img_cv = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
+        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
+        
+        # 2. Read Mask
+        mask_file = request.files['mask']
+        mask_arr = np.frombuffer(mask_file.read(), np.uint8)
+        
+        # [CRITICAL FIX] Read "unchanged" to preserve the low values (1, 2, 3)
+        mask_cv = cv2.imdecode(mask_arr, cv2.IMREAD_UNCHANGED)
+        
+        # If mask is RGB/RGBA, convert to grayscale
+        if len(mask_cv.shape) > 2:
+            mask_cv = cv2.cvtColor(mask_cv, cv2.COLOR_BGR2GRAY)
+
+        # 3. Preprocess
+        img_h, img_w = img_cv.shape[:2]
+        img_resized = cv2.resize(img_cv, (512, 512))
+        
+        # Resize mask carefully (Nearest Neighbor preserves exact class IDs 0,1,2...)
+        mask_resized = cv2.resize(mask_cv, (512, 512), interpolation=cv2.INTER_NEAREST)
+        
+        # Normalize Image
+        img_tensor = (torch.tensor(img_resized).float() / 127.5) - 1.0
+        img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).to(DEVICE)
+        
+        # [CRITICAL FIX] Logic change: Check if pixel > 0, NOT > 127
+        # This converts your class indices (1, 2, 3...) into a binary 1.0
+        mask_tensor = (torch.tensor(mask_resized).float() > 0).float().unsqueeze(0).unsqueeze(0).to(DEVICE)
+        
+        # 4. Inference
+        with torch.no_grad():
+            output = model(img_tensor, mask_tensor)
+            
+        # 5. Post-process (Same as before)
+        output_np = output.squeeze().permute(1, 2, 0).cpu().numpy()
+        output_np = (output_np + 1.0) * 127.5
+        output_np = np.clip(output_np, 0, 255).astype(np.uint8)
+        output_final = cv2.resize(output_np, (img_w, img_h))
+        
+        return jsonify({'result': f"data:image/png;base64,{to_base64(output_final)}"})
+        
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)