Update app.py

app.py

@@ -266,20 +266,93 @@
 #     demo.launch()
 
 
-#############6th##################
+#############6th-worked-proc##################
+# import torch
+# import gradio as gr
+# import requests
+# import os
+# import numpy as np
+
+# # Hugging Face Model Repository
+# model_repo = "Mariam-Elz/CRM"
+
+# # Download Model Weights (Only CRM.pth to Save Memory)
+# model_path = "models/CRM.pth"
+# os.makedirs("models", exist_ok=True)
+
+# if not os.path.exists(model_path):
+#     url = f"https://huggingface.co/{model_repo}/resolve/main/CRM.pth"
+#     print(f"Downloading CRM.pth...")
+#     response = requests.get(url)
+#     with open(model_path, "wb") as f:
+#         f.write(response.content)
+
+# # Set Device (Use CPU to Reduce RAM Usage)
+# device = "cpu"
+
+# # Load Model Efficiently
+# def load_model():
+#     model = torch.load(model_path, map_location=device)
+#     if isinstance(model, torch.nn.Module):
+#         model.eval()  # Ensure model is in inference mode
+#     return model
+
+# # Load model only when needed (saves memory)
+# model = load_model()
+
+# # Define Inference Function with Memory Optimizations
+# def infer(image):
+#     """Process input image and return a reconstructed image."""
+#     with torch.no_grad():
+#         # Convert image to torch tensor & normalize (float16 to save RAM)
+#         image_tensor = torch.tensor(image, dtype=torch.float16).unsqueeze(0).permute(0, 3, 1, 2) / 255.0
+#         image_tensor = image_tensor.to(device)
+
+#         # Model Inference
+#         output = model(image_tensor)
+
+#         # Convert back to numpy image format
+#         output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0
+#         output_image = np.clip(output_image, 0, 255).astype(np.uint8)
+
+#         # Free Memory
+#         del image_tensor, output
+#         torch.cuda.empty_cache()
+
+#         return output_image
+
+# # Create Gradio UI
+# demo = gr.Interface(
+#     fn=infer,
+#     inputs=gr.Image(type="numpy"),
+#     outputs=gr.Image(type="numpy"),
+#     title="Optimized Convolutional Reconstruction Model",
+#     description="Upload an image to get the reconstructed output with reduced memory usage."
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+#############7tth################
 import torch
+import torch.nn as nn
 import gradio as gr
 import requests
 import os
+import torchvision.transforms as transforms
 import numpy as np
+from PIL import Image
 
 # Hugging Face Model Repository
 model_repo = "Mariam-Elz/CRM"
 
-# Download Model Weights (Only CRM.pth to Save Memory)
+# Model File Path
 model_path = "models/CRM.pth"
 os.makedirs("models", exist_ok=True)
 
+# Download model weights if not present
 if not os.path.exists(model_path):
     url = f"https://huggingface.co/{model_repo}/resolve/main/CRM.pth"
     print(f"Downloading CRM.pth...")
@@ -287,51 +360,81 @@ if not os.path.exists(model_path):
     with open(model_path, "wb") as f:
         f.write(response.content)
 
-# Set Device (Use CPU to Reduce RAM Usage)
-device = "cpu"
-
-# Load Model Efficiently
+# Set Device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Define Model Architecture (Replace with your actual model)
+class CRMModel(nn.Module):
+    def __init__(self):
+        super(CRMModel, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+        self.relu = nn.ReLU()
+    
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        return x
+
+# Load Model
 def load_model():
-    model = torch.load(model_path, map_location=device)
-    if isinstance(model, torch.nn.Module):
-        model.eval()  # Ensure model is in inference mode
+    print("Loading model...")
+    model = CRMModel()  # Use the correct architecture here
+    state_dict = torch.load(model_path, map_location=device)
+
+    if isinstance(state_dict, dict):  # Ensure it's a valid state_dict
+        model.load_state_dict(state_dict)
+    else:
+        raise ValueError("Error: The loaded state_dict is not in the correct format.")
+
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully!")
     return model
 
-# Load model only when needed (saves memory)
+# Load the model
 model = load_model()
 
-# Define Inference Function with Memory Optimizations
+# Define Inference Function
 def infer(image):
-    """Process input image and return a reconstructed image."""
-    with torch.no_grad():
-        # Convert image to torch tensor & normalize (float16 to save RAM)
-        image_tensor = torch.tensor(image, dtype=torch.float16).unsqueeze(0).permute(0, 3, 1, 2) / 255.0
-        image_tensor = image_tensor.to(device)
-
-        # Model Inference
-        output = model(image_tensor)
-
-        # Convert back to numpy image format
-        output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0
-        output_image = np.clip(output_image, 0, 255).astype(np.uint8)
-
-        # Free Memory
-        del image_tensor, output
-        torch.cuda.empty_cache()
-
-        return output_image
+    """Process input image and return a reconstructed 3D output."""
+    try:
+        print("Preprocessing image...")
+
+        # Convert image to PyTorch tensor & normalize
+        transform = transforms.Compose([
+            transforms.Resize((256, 256)),  # Resize to fit model input
+            transforms.ToTensor(),  # Converts to tensor (C, H, W)
+            transforms.Normalize(mean=[0.5], std=[0.5]),  # Normalize
+        ])
+        image_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension
+
+        print("Running inference...")
+        with torch.no_grad():
+            output = model(image_tensor)  # Forward pass
+
+        # Ensure output is a valid tensor
+        if isinstance(output, torch.Tensor):
+            output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
+            output_image = np.clip(output_image * 255.0, 0, 255).astype(np.uint8)
+            print("Inference complete! Returning output.")
+            return output_image
+        else:
+            print("Error: Model output is not a tensor.")
+            return None
+
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return None
 
 # Create Gradio UI
 demo = gr.Interface(
     fn=infer,
-    inputs=gr.Image(type="numpy"),
+    inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="numpy"),
-    title="Optimized Convolutional Reconstruction Model",
-    description="Upload an image to get the reconstructed output with reduced memory usage."
+    title="Convolutional Reconstruction Model",
+    description="Upload an image to get the reconstructed output."
 )
 
 if __name__ == "__main__":
     demo.launch()
-
-
-