Serkan Ozturk commited on
Commit
687e8c3
·
1 Parent(s): 0f72f6f
Files changed (2) hide show
  1. handler.py +6 -16
  2. serkan.py +38 -0
handler.py CHANGED
@@ -2,11 +2,12 @@ from typing import Dict, List, Any
2
  import torch
3
  import numpy as np
4
  import torch.nn.functional as F
5
-
6
  class EndpointHandler():
7
- def __init__(self, path="FiveK.pth"):
8
  # load the optimized model
9
- self.model = torch.load(path)
 
10
 
11
 
12
 
@@ -25,19 +26,8 @@ class EndpointHandler():
25
  img = inputs["image"]
26
 
27
  # Load the image
28
- img = np.float32(img) / 255. # Load and normalize the image
29
-
30
- # Convert to torch tensor and add batch dimension
31
- img_tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
32
-
33
- # Padding if necessary (to make image dimensions multiples of 4)
34
- b, c, h, w = img_tensor.shape
35
- factor = 4 # Assuming factor is 4, based on the code context
36
- H, W = ((h + factor) // factor) * factor, ((w + factor) // factor) * factor
37
- padh = H - h if h % factor != 0 else 0
38
- padw = W - w if w % factor != 0 else 0
39
- img_tensor = F.pad(img_tensor, (0, padw, 0, padh), 'reflect')
40
 
41
- restored = self.model(img_tensor)
42
  # postprocess the prediction
43
  return "OKAY"
 
2
  import torch
3
  import numpy as np
4
  import torch.nn.functional as F
5
+ from serkan import SimpleUpscaleModel
6
  class EndpointHandler():
7
+ def __init__(self, path="model_weights.pth"):
8
  # load the optimized model
9
+ self.model = SimpleUpscaleModel()
10
+ self.model.load_state_dict(torch.load(self.path))
11
 
12
 
13
 
 
26
  img = inputs["image"]
27
 
28
  # Load the image
29
+ img = np.float32(img)
 
 
 
 
 
 
 
 
 
 
 
30
 
31
+ upscaled = self.model(img)
32
  # postprocess the prediction
33
  return "OKAY"
serkan.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ class SimpleUpscaleModel(nn.Module):
5
+ def __init__(self, scale_factor=2):
6
+ """
7
+ A simple model for upscaling inputs using bilinear interpolation.
8
+
9
+ Args:
10
+ scale_factor (int): The factor by which to upscale the input.
11
+ """
12
+ super(SimpleUpscaleModel, self).__init__()
13
+
14
+ # Upsampling layer
15
+ self.upsample = nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=True)
16
+
17
+ def forward(self, x):
18
+ """
19
+ Forward pass of the network.
20
+
21
+ Args:
22
+ x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).
23
+
24
+ Returns:
25
+ torch.Tensor: Upscaled output tensor.
26
+ """
27
+ return self.upsample(x)
28
+
29
+ if __name__ == "__main__":
30
+ # Create the model
31
+ scale_factor = 2
32
+ model = SimpleUpscaleModel(scale_factor=scale_factor)
33
+
34
+ # Save the model
35
+ model_path = "model_weights.pth"
36
+ torch.save(model.state_dict(), model_path)
37
+
38
+ print(f"Model saved to {model_path}")