pytorch-retinanet_visualize.py

Pytorch实现的Retinanet 源码解读，代码参考.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import torchvision
import time
import os
import copy
import pdb
# Test file
import time
import argparse

import sys
import cv2

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, models, transforms

from dataloader import CocoDataset, CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, UnNormalizer, Normalizer

# Make sure the torch version is "0.4.*"
assert torch.__version__.split('.')[1] == '4'

# Check whether cuda is available
print('CUDA available: {}'.format(torch.cuda.is_available()))


def main(args=None):
	# Add --help information
	parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')

	parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.')
	parser.add_argument('--coco_path', help='Path to COCO directory')
	parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
	parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')

	parser.add_argument('--model', help='Path to model (.pt) file.')

	parser = parser.parse_args(args) 

	# Choose the dataset type
	if parser.dataset == 'coco':
		dataset_val = CocoDataset(parser.coco_path, set_name='val2017', transform=transforms.Compose([Normalizer(), Resizer()]))
	elif parser.dataset == 'csv':
		dataset_val = CSVDataset(train_file=parser.csv_train, class_list=parser.csv_classes, transform=transforms.Compose([Normalizer(), Resizer()]))
	else:
		raise ValueError('Dataset type not understood (must be csv or coco), exiting.')

	sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
	print(sampler_val)
	dataloader_val = DataLoader(dataset_val, num_workers=1, collate_fn=collater, batch_sampler=sampler_val)
	
	# Load pytorch model according to the corresponding model directory
	retinanet = torch.load(parser.model)

	use_gpu = True

	if use_gpu:
		retinanet = retinanet.cuda()

	# Convert the model to test mode as dropout and bn in the process of test are different from those in training. 
	retinanet.eval()

	# tensor = tensor*mean + std
	# Create the object
	unnormalize = UnNormalizer()

	def draw_caption(image, box, caption):
		print('Box is: {}'.format(box))
		b = np.array(box).astype(int) 
		print("B is: {}".format(b))
		cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
		cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)

	# print(list(dataloader_val))
	# return
	for idx, data in enumerate(dataloader_val):
		# print("idx and data are {} and {}\n".format(idx, data))

		# If we do not need bp (inference or testing), use no_grad() to save memory and do not store gradient data
		with torch.no_grad():
			st = time.time()
			# print('img dimension is {}'.format(data['img'].size())) # image is presented in this form (nums, channel, height, width)
			scores, classification, transformed_anchors = retinanet(data['img'].cuda().float())  # These parameters are variated
			# Scores are stored in a decended order, just print the stores which are greater than the threshold.
			# print('scores are {}, classification are {}, transfermed_anchors are {}'.format(scores, classification.size(), transformed_anchors.size()))
			print('Elapsed time: {}'.format(time.time()-st))
			idxs = np.where(scores>0.5) # choose scores>0.5 to present

			# image normalization, set to [0,255]
			img = np.array(255 * unnormalize(data['img'][0, :, :, :])).copy()
			img[img<0] = 0
			img[img>255] = 255  # the length of img is 3

			# print(len(img))
			img = np.transpose(img, (1, 2, 0))
			# print(len(img)) # the transposed img length is 640

			img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2RGB)

			# print the object bounding boxes with labels
			for j in range(idxs[0].shape[0]):
				bbox = transformed_anchors[idxs[0][j], :]
				x1 = int(bbox[0])
				y1 = int(bbox[1])
				x2 = int(bbox[2])
				y2 = int(bbox[3])
				label_name = dataset_val.labels[int(classification[idxs[0][j]])]
				draw_caption(img, (x1, y1, x2, y2), label_name)

				cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=2)
				print(label_name)

			cv2.imshow('img', img)
			# cv2.waitKey(100)
			cv2.waitKey(0)



if __name__ == '__main__':
 main()