Update app.py

app.py

@@ -8,14 +8,19 @@ from PIL import Image
 # Load your PyTorch model
 model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()
 
+
 # Define a function to preprocess the DICOM
-def preprocess_dicom(dicom_path):
+def preprocess_dicom_to_image(dicom_path):
     # Load DICOM file
     dicom = pydicom.dcmread(dicom_path)
     image = dicom.pixel_array  # Extract image data
     # Normalize to [0, 1] and convert to PIL Image for transforms
     image = (image - np.min(image)) / (np.max(image) - np.min(image))
-    image = Image.fromarray((image * 255).astype(np.uint8))
+    return Image.fromarray((image * 255).astype(np.uint8))
+
+# Define a function to preprocess the DICOM
+def preprocess_dicom(dicom_path):
+    image = preprocess_dicom_to_image(dicom_path)
     # Apply transforms
     transform = transforms.Compose([
         transforms.Resize((224, 224)),  # Resize to model's input size
@@ -36,9 +41,13 @@ def predict_dicom(dicom_file):
     output_image = Image.fromarray(output_image.astype(np.uint8))
     return output_image
 
+# Prediction function
+def predict_dicom2(dicom_file):
+    return preprocess_dicom_to_image(dicom_file)
+
 # Create Gradio interface
 interface = gr.Interface(
-    fn=predict_dicom,
+    fn=predict_dicom2,
     inputs=gr.File(label="Upload DICOM File"),
     outputs="image",
     title="DICOM Image Prediction"