Update app.py

app.py

@@ -6,50 +6,38 @@ from facenet_pytorch import MTCNN
 from torchvision.transforms.functional import to_pil_image
 
 # Function to load the ViT model and MTCNN
-def load_model_and_mtcnn(model_path):
+def load_model(model_path):
     model = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
-    mtcnn = MTCNN(keep_all=True, device=device)
-    return model, device, mtcnn
+    return model, device
 
-def detect_and_process_skin(image_bytes):
-    """Detects faces in an image, crops the skin region, and returns it as an image object."""
-    # Load image from bytes
-    img = Image.open(io.BytesIO(image_bytes))
-    img_np = np.array(img)
-    img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
+# Initialize MTCNN for face detection
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+mtcnn = MTCNN(keep_all=True, device=device)
 
-    # Detect faces in the image
-    detections = mtcnn.detect_faces(img_rgb)
+# Function to preprocess the image using MTCNN for face detection
+def preprocess_image(image, device):
+    # Convert PIL image to OpenCV format
+    open_cv_image = np.array(image)
+    # Convert RGB to BGR for OpenCV
+    open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
 
-    # Check if any faces were detected
-    if detections:
-        x, y, width, height = detections[0]['box']
+    # Convert OpenCV image back to PIL Image for MTCNN
+    pil_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))
 
-        # Crop the face region
-        face_img_np = img_np[y:y+height, x:x+width]
-
-        # Convert to PIL Image for return
-        pil_img = Image.fromarray(face_img_np)
-        return pil_img
+    # Use MTCNN to detect faces
+    boxes, _ = mtcnn.detect(pil_image)
+    if boxes is not None:
+        # Crop the first detected face (for simplicity)
+        box = boxes[0].astype(int)
+        cropped_face = open_cv_image[box[1]:box[3], box[0]:box[2]]
+        # Convert cropped face back to PIL for further processing
+        processed_image = Image.fromarray(cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB))
     else:
-        # Return original image if no face was detected
-        return img
+        processed_image = image  # Use the original image if no face is detected
 
-# Function to preprocess the image and return both the tensor and the final PIL image for display
-def preprocess_image(image, mtcnn, device):
-    processed_image = image  # Initialize with the original image
-    try:
-        # Directly call mtcnn with the image to get cropped faces
-        cropped_faces = mtcnn(image)
-        if cropped_faces is not None and len(cropped_faces) > 0:
-            # Convert the first detected face tensor back to PIL Image for further processing
-            processed_image = to_pil_image(cropped_faces[0].cpu(),mode='BGR;16')
-    except Exception as e:
-        st.write(f"Exception in face detection: {e}")
-        processed_image = image
-    
+    # Transform image for model
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -80,14 +68,9 @@ uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption='Uploaded Image', use_column_width=True)
-    image1 = image.getvalue()
-    image_ten = detect_and_process_skin(image1)
     image_tensor, final_image = preprocess_image(image, mtcnn, device)
-    predicted_class, probabilities = predict(image_ten, model, device)
+    predicted_class, probabilities = predict(image_tensor, model, device)
     
     st.write(f"Predicted class: {predicted_class.item()}")
     # Display the final processed image
-   # st.image(final_image, caption='Processed Image', use_column_width=True)
-    img_bytes = io.BytesIO()
-    detect_and_process_skin(image1.getvalue()).save(img_bytes, format='JPEG')
-    st.image(img_bytes.getvalue(), width=250, caption="Processed Image")
+    st.image(final_image, caption='Processed Image', use_column_width=True)