Skip to content

Commit 8af8f6b

Browse files
author
Anastasia
committed
Added PyTorch to TensorFlow model conversion
1 parent 638d02d commit 8af8f6b

File tree

7 files changed

+1220
-0
lines changed

7 files changed

+1220
-0
lines changed
Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import cv2
2+
import numpy as np
3+
import tensorflow as tf
4+
import torch
5+
from albumentations import (
6+
Compose,
7+
Normalize,
8+
)
9+
from pytorch2keras.converter import pytorch_to_keras
10+
from torch.autograd import Variable
11+
12+
from PyTorchFullyConvolutionalResnet18 import FullyConvolutionalResnet18
13+
14+
15+
def converted_fully_convolutional_resnet18(
16+
input_tensor, pretrained_resnet=True,
17+
):
18+
# define input tensor
19+
input_var = Variable(torch.FloatTensor(input_tensor))
20+
21+
# get PyTorch ResNet18 model
22+
model_to_transfer = FullyConvolutionalResnet18(pretrained=pretrained_resnet)
23+
model_to_transfer.eval()
24+
25+
# convert PyTorch model to Keras
26+
model = pytorch_to_keras(
27+
model_to_transfer,
28+
input_var,
29+
[input_var.shape[-3:]],
30+
change_ordering=True,
31+
verbose=False,
32+
name_policy="keep",
33+
)
34+
35+
return model
36+
37+
38+
if __name__ == "__main__":
39+
# read ImageNet class ids to a list of labels
40+
with open("imagenet_classes.txt") as f:
41+
labels = [line.strip() for line in f.readlines()]
42+
43+
# read image
44+
original_image = cv2.imread("camel.jpg")
45+
46+
# convert original image to RGB format
47+
image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
48+
49+
# transform input image:
50+
transform = Compose(
51+
[
52+
Normalize(
53+
# subtract mean
54+
mean=(0.485, 0.456, 0.406),
55+
# divide by standard deviation
56+
std=(0.229, 0.224, 0.225),
57+
),
58+
],
59+
)
60+
# apply image transformations, (725, 1920, 3)
61+
image = transform(image=image)["image"]
62+
63+
# NHWC: (1, 725, 1920, 3)
64+
predict_image = tf.expand_dims(image, 0)
65+
# NCHW: (1, 3, 725, 1920)
66+
image = np.transpose(tf.expand_dims(image, 0).numpy(), [0, 3, 1, 2])
67+
68+
# get transferred torch ResNet18 with pre-trained ImageNet weights
69+
model = converted_fully_convolutional_resnet18(
70+
input_tensor=image, pretrained_resnet=True,
71+
)
72+
73+
# Perform inference.
74+
# Instead of a 1×1000 vector, we will get a
75+
# 1×1000×n×m output ( i.e. a probability map
76+
# of size n × m for each 1000 class,
77+
# where n and m depend on the size of the image).
78+
preds = model.predict(predict_image)
79+
# NHWC: (1, 3, 8, 1000) back to NCHW: (1, 1000, 3, 8)
80+
preds = tf.transpose(preds, (0, 3, 1, 2))
81+
preds = tf.nn.softmax(preds, axis=1)
82+
print("Response map shape : ", preds.shape)
83+
84+
# find the class with the maximum score in the n x m output map
85+
pred = tf.math.reduce_max(preds, axis=1)
86+
class_idx = tf.math.argmax(preds, axis=1)
87+
88+
row_max = tf.math.reduce_max(pred, axis=1)
89+
row_idx = tf.math.argmax(pred, axis=1)
90+
91+
col_idx = tf.math.argmax(row_max, axis=1)
92+
93+
predicted_class = tf.gather_nd(
94+
class_idx, (0, tf.gather_nd(row_idx, (0, col_idx[0])), col_idx[0]),
95+
)
96+
97+
# print top predicted class
98+
print("Predicted Class : ", labels[predicted_class], predicted_class)
99+
100+
# find the n × m score map for the predicted class
101+
score_map = tf.expand_dims(preds[0, predicted_class, :, :], 0).numpy()
102+
score_map = score_map[0]
103+
104+
# resize score map to the original image size
105+
score_map = cv2.resize(
106+
score_map, (original_image.shape[1], original_image.shape[0]),
107+
)
108+
109+
# binarize score map
110+
_, score_map_for_contours = cv2.threshold(
111+
score_map, 0.25, 1, type=cv2.THRESH_BINARY,
112+
)
113+
114+
score_map_for_contours = score_map_for_contours.astype(np.uint8).copy()
115+
116+
# Find the contour of the binary blob
117+
contours, _ = cv2.findContours(
118+
score_map_for_contours, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE,
119+
)
120+
121+
# find bounding box around the object.
122+
rect = cv2.boundingRect(contours[0])
123+
124+
# apply score map as a mask to original image
125+
score_map = score_map - np.min(score_map[:])
126+
score_map = score_map / np.max(score_map[:])
127+
128+
score_map = cv2.cvtColor(score_map, cv2.COLOR_GRAY2BGR)
129+
masked_image = (original_image * score_map).astype(np.uint8)
130+
131+
# display bounding box
132+
cv2.rectangle(
133+
masked_image, rect[:2], (rect[0] + rect[2], rect[1] + rect[3]), (0, 0, 255), 2,
134+
)
135+
136+
# display images
137+
cv2.imshow("Original Image", original_image)
138+
cv2.imshow("scaled_score_map", score_map)
139+
cv2.imshow("activations_and_bbox", masked_image)
140+
cv2.waitKey(0)
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import torch
2+
import torch.nn as nn
3+
from torch.hub import load_state_dict_from_url
4+
from torchvision import models
5+
6+
7+
# Define the architecture by modifying resnet.
8+
# Original code is here
9+
# https://github.com/pytorch/vision/blob/b2e95657cd5f389e3973212ba7ddbdcc751a7878/torchvision/models/resnet.py
10+
class FullyConvolutionalResnet18(models.ResNet):
11+
def __init__(self, num_classes=1000, pretrained=False, **kwargs):
12+
13+
# Start with standard resnet18 defined here
14+
# https://github.com/pytorch/vision/blob/b2e95657cd5f389e3973212ba7ddbdcc751a7878/torchvision/models/resnet.py
15+
super().__init__(
16+
block=models.resnet.BasicBlock,
17+
layers=[2, 2, 2, 2],
18+
num_classes=num_classes,
19+
**kwargs,
20+
)
21+
if pretrained:
22+
state_dict = load_state_dict_from_url(
23+
models.resnet.model_urls["resnet18"], progress=True,
24+
)
25+
self.load_state_dict(state_dict)
26+
27+
# Replace AdaptiveAvgPool2d with standard AvgPool2d
28+
# https://github.com/pytorch/vision/blob/b2e95657cd5f389e3973212ba7ddbdcc751a7878/torchvision/models/resnet.py#L153-L154
29+
self.avgpool = nn.AvgPool2d((7, 7))
30+
31+
# Add final Convolution Layer.
32+
self.last_conv = torch.nn.Conv2d(
33+
in_channels=self.fc.in_features, out_channels=num_classes, kernel_size=1,
34+
)
35+
self.last_conv.weight.data.copy_(
36+
self.fc.weight.data.view(*self.fc.weight.data.shape, 1, 1),
37+
)
38+
self.last_conv.bias.data.copy_(self.fc.bias.data)
39+
40+
# Reimplementing forward pass.
41+
# Replacing the following code
42+
# https://github.com/pytorch/vision/blob/b2e95657cd5f389e3973212ba7ddbdcc751a7878/torchvision/models/resnet.py#L197-L213
43+
def _forward_impl(self, x):
44+
# Standard forward for resnet18
45+
x = self.conv1(x)
46+
x = self.bn1(x)
47+
x = self.relu(x)
48+
x = self.maxpool(x)
49+
50+
x = self.layer1(x)
51+
x = self.layer2(x)
52+
x = self.layer3(x)
53+
x = self.layer4(x)
54+
x = self.avgpool(x)
55+
56+
# Notice, there is no forward pass
57+
# through the original fully connected layer.
58+
# Instead, we forward pass through the last conv layer
59+
x = self.last_conv(x)
60+
return x
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
This contains the code for **PyTorch to Tensorflow Model Conversion**. For more information - visit
2+
[**PyTorch to Tensorflow Model Conversion**](https://www.learnopencv.com/pytorch-to-tensorflow-model-conversion/)
3+
4+
# AI Courses by OpenCV
5+
6+
Want to become an expert in AI? [AI Courses by OpenCV](https://opencv.org/courses/) is a great place to start.
7+
8+
<a href="https://opencv.org/courses/">
9+
<p align="center">
10+
<img src="https://www.learnopencv.com/wp-content/uploads/2020/04/AI-Courses-By-OpenCV-Github.png">
11+
</p>
12+
</a>
338 KB
Loading

0 commit comments

Comments
 (0)