Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image, ImageOps | |
| import matplotlib.pyplot as plt | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchvision import transforms | |
| import os | |
| import time | |
| import io | |
| import base64 | |
| class Net2(nn.Module): | |
| def __init__(self): | |
| super(Net2, self).__init__() | |
| self.conv1 = nn.Conv2d(3, 64, 3, padding=1) | |
| self.bn1 = nn.BatchNorm2d(64) | |
| self.pool1 = nn.MaxPool2d(2, 2) | |
| self.dropout1 = nn.Dropout(0.25) | |
| self.conv2 = nn.Conv2d(64, 64, 3, padding=1) | |
| self.bn2 = nn.BatchNorm2d(64) | |
| self.pool2 = nn.MaxPool2d(2, 2) | |
| self.dropout2 = nn.Dropout(0.25) | |
| self.conv3 = nn.Conv2d(64, 64, 3, padding=1) | |
| self.bn3 = nn.BatchNorm2d(64) | |
| self.pool3 = nn.MaxPool2d(2, 2) | |
| self.dropout3 = nn.Dropout(0.25) | |
| self.conv4 = nn.Conv2d(64, 64, 3, padding=1) | |
| self.bn4 = nn.BatchNorm2d(64) | |
| self.pool4 = nn.MaxPool2d(2, 2) | |
| self.dropout4 = nn.Dropout(0.25) | |
| self.flatten = nn.Flatten() | |
| self.fc1 = nn.Linear(64 * 5 * 5, 200) | |
| self.fc2 = nn.Linear(200, 150) | |
| self.fc3 = nn.Linear(150, 2) | |
| def forward(self, x): | |
| x = F.relu(self.bn1(self.conv1(x))) | |
| x = self.pool1(x) | |
| x = self.dropout1(x) | |
| x = F.relu(self.bn2(self.conv2(x))) | |
| x = self.pool2(x) | |
| x = self.dropout2(x) | |
| x = F.relu(self.bn3(self.conv3(x))) | |
| x = self.pool3(x) | |
| x = self.dropout3(x) | |
| x = F.relu(self.bn4(self.conv4(x))) | |
| x = self.pool4(x) | |
| x = self.dropout4(x) | |
| x = self.flatten(x) | |
| x = F.relu(self.fc1(x)) | |
| x = F.relu(self.fc2(x)) | |
| x = F.softmax(self.fc3(x), dim=1) | |
| return x | |
| model2 = None | |
| model2_path = "model4.pth" | |
| if os.path.exists(model2_path): | |
| state_dict = torch.load(model2_path, map_location=torch.device('cpu')) | |
| new_state_dict = {} | |
| for key, value in state_dict.items(): | |
| new_key = key.replace("module.", "") | |
| new_state_dict[new_key] = value | |
| model = Net2() | |
| model.load_state_dict(new_state_dict) | |
| model.eval() | |
| else: | |
| print("Model file not found at", model2_path) | |
| def process_image(input_image): | |
| image = Image.open(io.BytesIO(input_image)).convert("RGB") | |
| start_time = time.time() | |
| heatmap = scanmap(np.array(image), model) | |
| elapsed_time = time.time() - start_time | |
| heatmap_img = Image.fromarray(np.uint8(plt.cm.hot(heatmap) * 255)).convert('RGB') | |
| heatmap_img = heatmap_img.resize(image.size) | |
| return image, heatmap_img, int(elapsed_time) | |
| def scanmap(image_np, model): | |
| image_np = image_np.astype(np.float32) / 255.0 | |
| window_size = (80, 80) | |
| stride = 10 | |
| height, width, channels = image_np.shape | |
| probabilities_map = [] | |
| for y in range(0, height - window_size[1] + 1, stride): | |
| row_probabilities = [] | |
| for x in range(0, width - window_size[0] + 1, stride): | |
| cropped_window = image_np[y:y + window_size[1], x:x + window_size[0]] | |
| cropped_window_torch = transforms.ToTensor()(cropped_window).unsqueeze(0) | |
| with torch.no_grad(): | |
| probabilities = model(cropped_window_torch) | |
| row_probabilities.append(probabilities[0, 1].item()) | |
| probabilities_map.append(row_probabilities) | |
| probabilities_map = np.array(probabilities_map) | |
| return probabilities_map | |
| def gradio_process_image(input_image): | |
| original, heatmap, elapsed_time = process_image(input_image.read()) | |
| return original, heatmap, f"Elapsed Time (seconds): {elapsed_time}" | |
| inputs = gr.Image(label="Upload Image") | |
| outputs = [ | |
| gr.Image(label="Original Image"), | |
| gr.Image(label="Heatmap"), | |
| gr.Textbox(label="Elapsed Time") | |
| ] | |
| iface = gr.Interface(fn=gradio_process_image, inputs=inputs, outputs=outputs) | |
| iface.launch() | |