Commit 09995c03 authored by Yuxin Wu's avatar Yuxin Wu

[FasterRCNN] Let viz not depend on coco

parent 2766bb5a
......@@ -15,11 +15,13 @@ from tensorpack.utils.timer import timed_operation
from tensorpack.utils.argtools import log_once
from pycocotools.coco import COCO
import config
__all__ = ['COCODetection', 'COCOMeta']
COCO_NUM_CATEGORY = 80
config.NUM_CLASS = COCO_NUM_CATEGORY + 1
class _COCOMeta(object):
......@@ -49,6 +51,7 @@ class _COCOMeta(object):
v: i + 1 for i, v in enumerate(cat_ids)}
self.class_id_to_category_id = {
v: k for k, v in self.category_id_to_class_id.items()}
config.CLASS_NAMES = self.class_names
COCOMeta = _COCOMeta()
......
......@@ -9,6 +9,7 @@ BASEDIR = '/path/to/your/COCO/DIR'
TRAIN_DATASET = ['train2014', 'valminusminival2014']
VAL_DATASET = 'minival2014' # only support evaluation on one dataset
NUM_CLASS = 81
CLASS_NAMES = [] # NUM_CLASS strings
# basemodel ----------------------
RESNET_NUM_BLOCK = [3, 4, 6, 3] # resnet50
......
......@@ -12,6 +12,7 @@ This file is modified from
https://github.com/tensorflow/models/blob/master/object_detection/core/box_list_ops.py
"""
@under_name_scope()
def area(boxes):
"""
......@@ -73,7 +74,7 @@ def get_iou_callable():
"""
Get a pairwise box iou callable.
"""
with tf.device('/cpu:0'):
with tf.Graph().as_default(), tf.device('/cpu:0'):
A = tf.placeholder(tf.float32, shape=[None, 4])
B = tf.placeholder(tf.float32, shape=[None, 4])
iou = pairwise_iou(A, B)
......
......@@ -8,8 +8,8 @@ import numpy as np
from tensorpack.utils import viz
from tensorpack.utils.palette import PALETTE_RGB
from coco import COCOMeta
from utils.box_ops import get_iou_callable
import config
def draw_annotation(img, boxes, klass, is_crowd=None):
......@@ -18,13 +18,13 @@ def draw_annotation(img, boxes, klass, is_crowd=None):
if is_crowd is not None:
assert len(boxes) == len(is_crowd)
for cls, crd in zip(klass, is_crowd):
clsname = COCOMeta.class_names[cls]
clsname = config.CLASS_NAMES[cls]
if crd == 1:
clsname += ';Crowd'
labels.append(clsname)
else:
for cls in klass:
labels.append(COCOMeta.class_names[cls])
labels.append(config.CLASS_NAMES[cls])
img = viz.draw_boxes(img, boxes, labels)
return img
......@@ -59,7 +59,7 @@ def draw_predictions(img, boxes, scores):
return img
labels = scores.argmax(axis=1)
scores = scores.max(axis=1)
tags = ["{},{:.2f}".format(COCOMeta.class_names[lb], score) for lb, score in zip(labels, scores)]
tags = ["{},{:.2f}".format(config.CLASS_NAMES[lb], score) for lb, score in zip(labels, scores)]
return viz.draw_boxes(img, boxes, tags)
......@@ -74,7 +74,7 @@ def draw_final_outputs(img, results):
tags = []
for label, _, score in results:
tags.append(
"{},{:.2f}".format(COCOMeta.class_names[label], score))
"{},{:.2f}".format(config.CLASS_NAMES[label], score))
boxes = np.asarray([x.box for x in results])
return viz.draw_boxes(img, boxes, tags)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment