first commit

app.py

@@ -0,0 +1,16 @@
+import gradio as gr
+from predict import predict_one_image
+
+input_image = [
+    gr.components.Image(type='filepath',label='Input Image')
+]
+label = gr.outputs.Label()
+examples = ['10061.jpg','10013.jpg']
+gr.Interface(
+    fn=predict_one_image,
+    inputs=input_image,
+    outputs=label,
+    title="Ai Image Generated Detector app",
+    examples=examples,
+    cache_examples=False,
+).launch(share=True)

model.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+import timm
+from torch.nn.parameter import Parameter
+
+class Backbone(nn.Module) :
+    def __init__(self,name,pretrained) :
+        super(Backbone,self).__init__()
+        self.net = timm.create_model(name,pretrained=pretrained)
+        self.out_features = self.net.get_classifier().in_features
+    def forward(self,x) :
+        x = self.net.forward_features(x)
+        return x
+
+class CustomModel(nn.Module) :
+    def __init__(self) :
+        super(CustomModel,self).__init__()
+        self.backbone = Backbone("tf_efficientnetv2_b0",False)
+        self.pooling = nn.AdaptiveAvgPool2d(1)
+        self.head = nn.Linear(self.backbone.out_features,1)
+    def forward(self,x) :
+        x = self.backbone(x)
+        x = self.pooling(x).squeeze()
+        target = self.head(x)
+        output = {}
+        output['label'] = target
+        return output

predict.py

@@ -0,0 +1,21 @@
+import cv2
+from utils import read_image,get_valid_augs
+import torch
+import torch.nn.functional as F
+from model import CustomModel
+
+CKPT = 'fold_0.pt'
+Targets = ['Not AI' 'AI Generated']
+def predict_one_image(path) :
+    image = read_image(path)
+    image = get_valid_augs()(image=image)['image']
+    image = torch.tensor(image,dtype=torch.float)
+    image = image.reshape((1,3,224,224))
+    model = CustomModel()
+    #loading ckpt
+    model.load_state_dict(torch.load(CKPT,map_location=torch.device('cpu')))
+    with torch.no_grad() :
+        outputs = model(image)
+        proba = F.sigmoid(outputs['label']).detach().numpy()[0]
+    return {'Not AI' : 1-float(proba),'AI' : float(proba)}#(proba>0.5)*1
+

utils.py

@@ -0,0 +1,22 @@
+import cv2
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+def read_image(path) :
+    img = cv2.imread(path)
+    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
+    return img
+
+def get_valid_augs() :
+    return  A.Compose([
+   A.Resize(height=224, width=224, always_apply=True, p=1),
+   A.Normalize(
+        mean = IMAGENET_DEFAULT_MEAN,
+        std  = IMAGENET_DEFAULT_STD,
+        max_pixel_value=255
+            ),
+   ToTensorV2(),
+])