Spaces:
Sleeping
Sleeping
File size: 1,844 Bytes
d3c0464 040a6f8 d3c0464 9a444fd d3c0464 9a444fd d3c0464 9a444fd d3c0464 9a444fd d3c0464 9a444fd 1e7bd9c d3c0464 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import torch
import numpy as np
import pydicom
import gradio as gr
from torchvision import transforms
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_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))
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
transforms.ToTensor(),
])
return transform(image).unsqueeze(0) # Add batch dimension
# Prediction function
def predict_dicom(dicom_file):
# Preprocess
input_tensor = preprocess_dicom(dicom_file.name)
# Inference
with torch.no_grad():
output = model(input_tensor)
# Convert output tensor to image (dummy example, replace as needed)
output_image = output.squeeze().numpy()
output_image = (output_image - np.min(output_image)) / (np.max(output_image) - np.min(output_image)) * 255
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_dicom2,
inputs=gr.File(label="Upload DICOM File"),
outputs="image",
title="DICOM Image Prediction"
)
# Launch the Gradio app
interface.launch()
|