Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import random | |
| import os | |
| from typing import Tuple , Dict | |
| import time | |
| import torch | |
| from PIL import Image | |
| import numpy as np | |
| from torchvision import transforms | |
| import torch.nn as nn | |
| from torch.nn.functional import relu | |
| import requests | |
| from io import BytesIO | |
| class GlobalAttention(nn.Module): | |
| def __init__(self, num_channels): | |
| super(GlobalAttention, self).__init__() | |
| self.attention = nn.Sequential( | |
| nn.Conv2d(num_channels, 1, kernel_size=1), | |
| nn.BatchNorm2d(1), | |
| nn.Sigmoid() | |
| ) | |
| def forward(self, x): | |
| attention_weights = self.attention(x) | |
| return x * attention_weights | |
| class ModelWithAttention(nn.Module): | |
| def __init__(self, num_characters): | |
| super(ModelWithAttention, self).__init__() | |
| self.conv1 = nn.Conv2d(1, 64, kernel_size=(3, 3), padding='same') | |
| self.bn1 = nn.BatchNorm2d(64) | |
| self.pool = nn.MaxPool2d(kernel_size=(2, 2)) | |
| self.conv2 = nn.Conv2d(64, 128, kernel_size=(3, 3), padding='same') | |
| self.bn2 = nn.BatchNorm2d(128) | |
| self.conv3 = nn.Conv2d(128, 256, kernel_size=(3, 3), padding='same') | |
| self.bn3 = nn.BatchNorm2d(256) | |
| self.conv4 = nn.Conv2d(256, 512, kernel_size=(3, 3), padding='same') | |
| self.bn4 = nn.BatchNorm2d(512) | |
| self.pool2 = nn.MaxPool2d(kernel_size=(1, 2)) | |
| self.attention = GlobalAttention(512) | |
| self.flatten = nn.Flatten() | |
| self.fc1 = nn.Linear(16384, 512) | |
| self.bn5 = nn.BatchNorm1d(512) | |
| self.dropout1 = nn.Dropout(0.5) | |
| self.fc2 = nn.Linear(512, 512) | |
| self.bn6 = nn.BatchNorm1d(512) | |
| self.dropout2 = nn.Dropout(0.75) | |
| self.output = nn.Linear(512, num_characters) | |
| self.sm = nn.Softmax() | |
| def forward(self, x): | |
| x = self.pool(relu(self.bn1(self.conv1(x)))) | |
| x = self.pool(relu(self.bn2(self.conv2(x)))) | |
| x = self.pool(relu(self.bn3(self.conv3(x)))) | |
| x = self.pool2(relu(self.bn4(self.conv4(x)))) | |
| # x = self.attention(x) | |
| x = self.flatten(x) | |
| x = relu(self.bn5(self.fc1(x))) | |
| x = self.dropout1(x) | |
| x = relu(self.bn6(self.fc2(x))) | |
| x = self.dropout2(x) | |
| x = self.output(x) | |
| x = self.sm(x) | |
| return x | |
| device = "cpu" | |
| path = "Captcha(Best).pt" | |
| from torchvision import transforms | |
| model = ModelWithAttention(10).to(device) | |
| model.load_state_dict(torch.load(path , map_location=torch.device('cpu'))) | |
| transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Grayscale(), | |
| transforms.Resize((64,64)) | |
| ]) | |
| def predict(img= None , link:str = None) -> str: | |
| sizes = [ | |
| [15,-5 , 15,27 ], | |
| [15,-5 , 28,40 ], | |
| [15,-5 , 41,53 ], | |
| [15,-5 , 53,65 ], | |
| [15,-5 , 66,78 ]] | |
| answer = "" | |
| if img != None: | |
| imgss = np.array((img)) | |
| model.eval() | |
| for size in (sizes): | |
| img = imgss[size[0]:size[1], size[2]:size[3]] | |
| img = Image.fromarray(img) | |
| img = transform(img) | |
| img = img.unsqueeze(0) | |
| answer += str((torch.argmax(model(img.to(device)))).numpy()) | |
| return answer , imgss | |
| if link != None: | |
| response = requests.get(str(link)) | |
| if response.status_code == 200: | |
| imgss = np.array(Image.open(BytesIO(response.content))) | |
| # print("Image downloaded and converted to numpy array successfully!") | |
| # print(imgss.shape) | |
| model.eval() | |
| for size in (sizes): | |
| img = imgss[size[0]:size[1], size[2]:size[3]] | |
| img = Image.fromarray(img) | |
| img = transform(img) | |
| img = img.unsqueeze(0) | |
| answer += str((torch.argmax(model(img.to(device)))).cpu().numpy()) | |
| return answer , imgss | |
| from pathlib import Path | |
| path = "example" | |
| list_path = [] | |
| list_paths = os.listdir(path) | |
| for i in list_paths: | |
| list_path.append(os.path.join(path , i)) | |
| # print(list_path) | |
| import gradio as gr | |
| title = "GIGA Captcha Solver" | |
| description = "This Model can solve persian numbers Captcha easly" | |
| article = "Created By A.M.Parviz <3" | |
| # Create the Gradio demo | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=[gr.Image(type="pil"), | |
| gr.Text()], | |
| outputs=[ | |
| gr.Label(num_top_classes=10, label="Predictions"), | |
| gr.Image() | |
| ], | |
| examples = [[img_path, ""] for img_path in list_path], | |
| title=title, | |
| description=description, | |
| article=article, | |
| ) | |
| demo.launch() | |
| # share=True) | |